Implement autodiff using intrinsics #142640

Sa4dUs · 2025-06-17T20:19:27Z

This PR aims to move autodiff logic to autodiff intrinsic. Allowing us to delete a great part of our frontend code and overall, simplify the compilation pipeline of autodiff functions.

workingjubilee · 2025-06-17T23:21:28Z

compiler/rustc_codegen_llvm/src/intrinsic.rs

+            _ if tcx.has_attr(def_id, sym::rustc_autodiff) => {
+                // NOTE(Sa4dUs): This is a hacky way to get the autodiff items
+                // so we can focus on the lowering of the intrinsic call
+
+                // `diff_items` is empty even when autodiff is enabled, and if we're here,
+                // it's because some function was marked as intrinsic and had the `rustc_autodiff` attr
+                let diff_items = tcx.collect_and_partition_mono_items(()).autodiff_items;
+
+                // this shouldn't happen?
+                if diff_items.is_empty() {
+                    bug!("no autodiff items found for {def_id:?}");
+                }
+
+                // TODO(Sa4dUs): generate the enzyme call itself, based on the logic in `builder.rs`
+
+                // Just gen the fallback body for now
+                return Err(ty::Instance::new_raw(def_id, instance.args));
+            }


So this is the part you got stuck on?

Kinda. I tried to use the collect_and_partition_mono_items just as a hacky way of getting autodiff_items, so i can focus on the declaration. Once that was working, I would focus on how to get that information in the best way possible.

Just for context, during yesterday's meeting, we came to the conclusion that, for the time being, it's acceptable to copy and paste the collector logic from within the provider's code, just to focus on the declaration, but still nice to have feedback since, in some moment, we'll need to get to AutoDiffItems in a decent way. collect_and_partition_mono_items(()).autodiff_items being [] in autodiff code seems weird, so it would be nice to know why it's happening, or if it is some kind of bug.

workingjubilee · 2025-06-17T23:37:58Z

@Sa4dUs Can you provide the motivation in the PR desc for context?

ZuseZ4 · 2025-07-15T23:10:33Z

tests/codegen/autodiff/batched.rs

-// CHECK-NEXT:   %19 = insertvalue [4 x float] %17, float %18, 3
-// CHECK-NEXT:   ret [4 x float] %19
-// CHECK-NEXT: }
+// CHECK-NEXT:   %4 = fadd fast float %"_2'ipl", %"_2'ipl"


Is this test currently passing on master? Also, can you compare the IR of the two rustc versions just before we run Enzyme? I wonder if we call enzyme differently, or if just llvm changed. (Ignore this if you're still working on the code and are not done yet).

yes, the test with the current fixes is passing at master. and the way we call enzyme is just the same, as i've based my work on the existing function for diff body codegen and removed unnecessary parts.

rustbot · 2025-07-27T16:46:45Z

r? @SparrowLii

rustbot has assigned @SparrowLii.
They will have a look at your PR within the next two weeks and either review your PR or reassign to another reviewer.

Use r? to explicitly pick a reviewer

rustbot · 2025-07-27T16:46:46Z

Some changes occurred in compiler/rustc_codegen_llvm/src/builder/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_codegen_gcc

cc @antoyo, @GuillaumeGomez

Some changes occurred in compiler/rustc_codegen_ssa/src/codegen_attrs.rs

cc @jdonszelmann

Some changes occurred to the intrinsics. Make sure the CTFE / Miri interpreter
gets adapted for the changes, if necessary.

cc @rust-lang/miri, @RalfJung, @oli-obk, @lcnr

These commits modify the Cargo.lock file. Unintentional changes to Cargo.lock can be introduced when switching branches and rebasing PRs.

If this was unintentional then you should revert the changes before this PR is merged.
Otherwise, you can ignore this comment.

Some changes occurred in compiler/rustc_codegen_ssa

cc @WaffleLapkin

Some changes occurred in compiler/rustc_builtin_macros/src/autodiff.rs

cc @ZuseZ4

Some changes occurred in compiler/rustc_monomorphize/src/partitioning/autodiff.rs

cc @ZuseZ4

triagebot.toml has been modified, there may have been changes to the review queue.

cc @davidtwco, @wesleywiser

library/core/src/intrinsics/mod.rs

SparrowLii · 2025-07-28T03:13:11Z

r? codegen

rustbot · 2025-07-30T07:35:59Z

The rustc-dev-guide subtree was changed. If this PR only touches the dev guide consider submitting a PR directly to rust-lang/rustc-dev-guide otherwise thank you for updating the dev guide with your changes.

cc @BoxyUwU, @jieyouxu, @Kobzol, @tshepang

bors · 2025-07-30T22:17:14Z

☔ The latest upstream changes (presumably #144692) made this pull request unmergeable. Please resolve the merge conflicts.

bors · 2025-08-04T21:29:12Z

☔ The latest upstream changes (presumably #144893) made this pull request unmergeable. Please resolve the merge conflicts.

bors · 2025-08-10T03:21:12Z

☔ The latest upstream changes (presumably #145146) made this pull request unmergeable. Please resolve the merge conflicts.

ZuseZ4

Looks mostly good to me after the last round, just some smaller comments before merging.

ZuseZ4 · 2025-08-13T18:25:51Z

compiler/rustc_monomorphize/src/collector/autodiff.rs

+// Here, we force both primal and diff function to be collected in
+// mono so this does not interfere in `enzyme_autodiff` intrinsics
+// codegen process. If they are unused, they will be removed later and
+// won't be present at LLVM-IR.


Can you clarify if later means on MIR level, or were they only absent because we usually run LLVM's O3 pipeline (which removes dead code).

it was because of O3 pipeline, fixing the comment

ZuseZ4 · 2025-08-13T18:32:07Z

compiler/rustc_span/src/symbol.rs

@@ -917,6 +917,7 @@ symbols! {
        enumerate_method,
        env,
        env_CFG_RELEASE: env!("CFG_RELEASE"),
+        enzyme_autodiff,


would it cause a name conflict to also just call this autodiff instead of enzyme_autodiff? The name is visible e.g. in the cargo expand output, and while Enzyme obviously does the heavy lifting atm., there are e.g. other backends to which autodiff could lower in the future, and they wouldn't use Enzyme.

no, there's no conflict, i'll rename it

ZuseZ4 · 2025-08-13T18:36:03Z

tests/codegen-llvm/autodiff/identical_fnc.rs

 // CHECK-NEXT:start:
 // CHECK-NOT:br
 // CHECK-NOT:ret
 // CHECK:; call identical_fnc::d_square
-// CHECK-NEXT:  call fastcc void @_ZN13identical_fnc8d_square17h4c364207a2f8e06dE(double %x.val, ptr noalias noundef nonnull align 8 dereferenceable(8) %dx1)
-// CHECK-NEXT:; call identical_fnc::d_square
-// CHECK-NEXT:  call fastcc void @_ZN13identical_fnc8d_square17h4c364207a2f8e06dE(double %x.val, ptr noalias noundef nonnull align 8 dereferenceable(8) %dx2)


It seems like autodiff again regressed, now we lack nonnull, even though rust references can't be null pointers. But we can debug that later, since it's "just" a performance question.

ZuseZ4 · 2025-08-13T18:41:34Z

compiler/rustc_codegen_llvm/src/intrinsic.rs

+    let callee_ty = instance.ty(tcx, bx.typing_env());
+
+    let sig = callee_ty.fn_sig(tcx);
+    let sig = tcx.normalize_erasing_late_bound_regions(bx.typing_env(), sig);


Can you leave an explanation on why this is the right thing to do here?

it's something i added in my first commits, when we didn't even have an autodiff intrinsic, and it doesn't seem to be necessary rn, so i'll just skip_binder() and things keep working as expected

it may have smth to do with the fact that the intrinsics where being generated dynamically instead of having a predefined one and smth would break, but can't remember exactly tbh

ZuseZ4 · 2025-08-13T18:44:06Z

compiler/rustc_codegen_llvm/src/intrinsic.rs

+            result
+        }
+
+        OperandValue::ZeroSized => bug!("unexpected ZeroSized argument in get_args_from_tuple"),


This is probably one of the locations we need to handle for #144025, right? I'd assume we just drop the corresponding activity.
In that case, can you leave a fixme?

this is matching the whole arg tuple passed to the intrinsic. just realized that tuple can be empty for silly functions like f() -> (), so i'm changing it so if the tuple is empty, it just return an empty vec

i added a fix for ZST in my other PR where i just drop the activity used for the ZST

ZuseZ4 · 2025-08-13T18:48:09Z

compiler/rustc_builtin_macros/src/autodiff.rs

-        }
-        body.stmts.push(ecx.stmt_semi(black_box_remaining_args));
+    // Generate turbofish expression from fn name and generics
+    // Given `foo` and `<A, B, C>` params, gen `foo::<A, B, C>`


Nit: Add a note that we use that when we pass the primal function to the autodiff intrinsic.

ZuseZ4 · 2025-08-13T18:51:27Z

compiler/rustc_builtin_macros/src/autodiff.rs

+
+        let segment = PathSegment {
+            ident,
+            id: ast::DUMMY_NODE_ID,


Can you add it as a test, so we don't regress it? Just a simple one that checks that this compiles. With a note saying that in the past we just test for the correct macro hygiene information .

ZuseZ4 · 2025-08-13T18:54:24Z

compiler/rustc_builtin_macros/src/autodiff.rs

@@ -566,202 +593,33 @@ mod llvm_enzyme {
    /// from optimizing any arguments away.
    fn gen_enzyme_body(


The comment above is outdated.

Actually, the function is down to just two lines, is only called once, and mostly just forwards input arguments, so just remove it.

ZuseZ4 · 2025-08-13T19:42:18Z

compiler/rustc_codegen_llvm/src/builder/autodiff.rs

 use crate::value::Value;
-use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};



Is there a reason why you moved this function to this file? It doesn't seem to use any LLVM functionality or types.
It's only called in the middle-end, https://github.com/rust-lang/rust/blob/master/compiler/rustc_monomorphize/src/partitioning/autodiff.rs#L126

not really, as i didn't need partitioning anymore, and i was calling that function directly in the intrinsic codegen, i guessed moving that to builder was the most "accurate" thing

wouldn't know what would be the correct place for it, as it's really not related with partitioning anymore, it's just an aux function used in codegen for getting the correct activities

also, didn't want to move to cg_llvm/intrinsic.rs file to not flood it with autodiff logic

ZuseZ4 · 2025-08-13T20:06:24Z

Also one follow-up. Now that the inline logic got rewritten, what about this PR that incorrectly tried to fix our logic (while actually breaking it?). #143633.
Can you also remove check from lto.rs?

Sa4dUs · 2025-08-14T08:20:50Z

actually, i can just remove the whole loop right?

bors · 2025-08-14T14:50:04Z

☔ The latest upstream changes (presumably #145388) made this pull request unmergeable. Please resolve the merge conflicts.

ZuseZ4 · 2025-08-14T18:07:07Z

tests/codegen-llvm/autodiff/macro_hygiene.rs

+
+#![feature(autodiff)]
+
+// CHECK: ;


Thanks for adding the test.
Don't we have an extra flag to check for something compiling, instead of an emtpy check?

rustbot added A-LLVM Area: Code generation parts specific to LLVM. Both correctness bugs and optimization-related issues. T-compiler Relevant to the compiler team, which will review and decide on the PR/issue. labels Jun 17, 2025

workingjubilee reviewed Jun 17, 2025

View reviewed changes

rustbot added the F-autodiff `#![feature(autodiff)]` label Jun 24, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch 2 times, most recently from bfb72fd to fad0b0c Compare July 1, 2025 10:53

This comment has been minimized.

Sign in to view

rustbot added the A-attributes Area: Attributes (`#[…]`, `#![…]`) label Jul 7, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from a72fd24 to 38b27a2 Compare July 7, 2025 17:30

This comment has been minimized.

Sign in to view

ZuseZ4 reviewed Jul 15, 2025

View reviewed changes

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 97f0985 to 43efce8 Compare July 25, 2025 17:38

rustbot added the A-meta Area: Issues & PRs about the rust-lang/rust repository itself label Jul 25, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 96f554c to d312894 Compare July 25, 2025 18:06

Sa4dUs marked this pull request as ready for review July 27, 2025 16:46

rustbot assigned SparrowLii Jul 27, 2025

rustbot added the S-waiting-on-review Status: Awaiting review from the assignee but also interested parties. label Jul 27, 2025

RalfJung reviewed Jul 27, 2025

View reviewed changes

library/core/src/intrinsics/mod.rs Outdated Show resolved Hide resolved

rustbot assigned saethlin and unassigned SparrowLii Jul 28, 2025

ZuseZ4 assigned ZuseZ4 and unassigned saethlin Jul 28, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from f1c7a40 to c370d96 Compare July 28, 2025 08:52

This comment has been minimized.

Sign in to view

ZuseZ4 mentioned this pull request Jul 30, 2025

clean up codegen fn attrs #144655

Merged

rust-cloud-vms bot force-pushed the ad-intrinsic branch from e24e7c4 to 126e73d Compare July 30, 2025 06:55

rustbot added the A-rustc-dev-guide Area: rustc-dev-guide label Jul 30, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 577cb5d to b7d7b73 Compare August 1, 2025 09:10

This comment has been minimized.

Sign in to view

rustbot added has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Aug 5, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 4cb3e3f to b9c9a69 Compare August 6, 2025 12:16

rustbot removed has-merge-commits PR has merge commits, merge with caution. S-waiting-on-author Status: This is awaiting some action (such as code changes or more information) from the author. labels Aug 6, 2025

rust-cloud-vms bot force-pushed the ad-intrinsic branch 2 times, most recently from 1eb9eee to 66fe728 Compare August 9, 2025 15:35

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 66fe728 to ff7594a Compare August 11, 2025 08:00

ZuseZ4 reviewed Aug 13, 2025

View reviewed changes

This comment has been minimized.

Sign in to view

rust-cloud-vms bot force-pushed the ad-intrinsic branch from fe9da82 to 9115b9f Compare August 14, 2025 15:47

Sa4dUs added 5 commits August 14, 2025 16:29

Basic implementation of autodiff intrinsic

5c63104

Complete functionality and general cleanup

250d77e

Remove inlining for autodiff handling

c9c1c17

Remove lto inline logic

e1d79b9

Update autodiff tests for the new intrinsics impl

4f1174d

rust-cloud-vms bot force-pushed the ad-intrinsic branch from 9115b9f to 4f1174d Compare August 14, 2025 16:36

ZuseZ4 reviewed Aug 14, 2025

View reviewed changes

		@@ -566,202 +593,33 @@ mod llvm_enzyme {
		/// from optimizing any arguments away.
		fn gen_enzyme_body(

		use crate::value::Value;
		use crate::{CodegenContext, LlvmCodegenBackend, ModuleLlvm, attributes, llvm};


		#![feature(autodiff)]

		// CHECK: ;

Implement autodiff using intrinsics #142640

Are you sure you want to change the base?

Implement autodiff using intrinsics #142640

Conversation

Sa4dUs commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Jun 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

workingjubilee commented Jun 17, 2025

Uh oh!

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

This comment has been minimized.

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

rustbot commented Jul 27, 2025

Uh oh!

rustbot commented Jul 27, 2025

Uh oh!

Uh oh!

SparrowLii commented Jul 28, 2025

Uh oh!

This comment has been minimized.

rustbot commented Jul 30, 2025

Uh oh!

bors commented Jul 30, 2025

Uh oh!

bors commented Aug 4, 2025

Uh oh!

This comment has been minimized.

bors commented Aug 10, 2025

Uh oh!

ZuseZ4 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Sa4dUs Aug 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ZuseZ4 commented Aug 13, 2025

Sa4dUs commented Jun 17, 2025 •

edited

Loading

Sa4dUs Jun 18, 2025 •

edited

Loading

Sa4dUs Aug 14, 2025 •

edited

Loading